import chainlit as cl
import asyncio
import logging
from src.workflow.agent_workflow import AgentState, agent_workflow
from src.core.knowledge_base import KnowledgeBase  # 用于文件上传处理

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化知识库
try:
    kb = KnowledgeBase()
    logger.info("知识库初始化成功")
except Exception as e:
    logger.error(f"知识库初始化失败: {e}")
    raise

@cl.on_chat_start
async def on_chat_start():
    # 添加文件上传组件
    files = await cl.AskFileMessage(
        content="请上传知识库文档（支持PDF、Word、TXT）",
        accept=["application/pdf",
                "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                "text/plain"],
        max_size_mb=20,
        timeout=180,
    ).send()
    
    if files:
        # 处理上传的文件
        success_count = 0
        for file in files:
            if kb.add_document(file.path):
                success_count += 1
        await cl.Message(content=f"已成功上传 {success_count}/{len(files)} 个文档到知识库").send()
    
    # 添加角色选择器
    roles = ["分析师", "风控师", "顾问", "全流程"]
    actions = [
        cl.Action(name=role, value=role, label=role) for role in roles
    ]
    await cl.Message(content="请选择角色：", actions=actions).send()
    
    # 初始化对话历史
    cl.user_session.set("history", [])
    await cl.Message(content="金融Agent已初始化，请开始对话").send()

@cl.on_message
async def on_message(message: cl.Message):
    # 获取用户选择的角色（默认为全流程）
    selected_role = cl.user_session.get("selected_role", "全流程")
    
    # 创建初始状态
    state = AgentState(message.content)
    state.selected_role = selected_role
    
    # 异步执行工作流
    result = await asyncio.to_thread(agent_workflow.invoke, state)
    
    # 更新对话历史
    history = cl.user_session.get("history")
    history.append({"user": message.content, "agent": result.final_advice})
    cl.user_session.set("history", history)
    
    # 发送最终建议
    await cl.Message(content=result.final_advice).send()

@cl.on_action
async def on_action(action: cl.Action):
    # 处理角色选择
    if action.name in ["分析师", "风控师", "顾问", "全流程"]:
        cl.user_session.set("selected_role", action.name)
        await cl.Message(content=f"已选择角色: {action.name}").send()
    
    # 处理历史对话查看
    if action.name == "view_history":
        history = cl.user_session.get("history", [])
        history_text = "\n".join([f"用户: {item['user']}\nAgent: {item['agent']}\n" for item in history])
        await cl.Message(content=f"对话历史:\n{history_text}").send()